//
// Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights
// Reserved. MIT License  (https://opensource.org/licenses/MIT)
//
/*
 * // 2022-2023 by zhaomingwork@qq.com
 */
// java FunasrWsClient
// usage:  FunasrWsClient [-h] [--port PORT] [--host HOST] [--audio_in AUDIO_IN] [--num_threads NUM_THREADS]
//                 [--chunk_size CHUNK_SIZE] [--chunk_interval CHUNK_INTERVAL] [--mode MODE]
package com.hjzd.ai.smartpractice.websocket;

import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.google.gson.LongSerializationPolicy;
import com.google.gson.reflect.TypeToken;
import com.hjzd.ai.smartpractice.config.WebClientConfig;
import com.hjzd.ai.smartpractice.constant.LLMRole;
import com.hjzd.ai.smartpractice.constant.ServiceConstant;
import com.hjzd.ai.smartpractice.entity.DTO.chat.ChatCompletionMessage;
import com.hjzd.ai.smartpractice.entity.DTO.chat.ChatCompletionRequest;
import com.hjzd.ai.smartpractice.entity.DTO.chat.ChatCompletionResponse;
import com.hjzd.ai.smartpractice.entity.DTO.chat.ChatResponseThread;
import com.hjzd.ai.smartpractice.service.RedisService;
import com.hjzd.ai.smartpractice.service.impl.ApiServiceImpl;
import org.apache.commons.lang3.StringUtils;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.drafts.Draft;
import org.java_websocket.handshake.ServerHandshake;
import org.json.simple.JSONArray;
import org.json.simple.JSONObject;
import org.json.simple.parser.JSONParser;
import org.json.simple.parser.ParseException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.client.reactive.ClientHttpConnector;
import org.springframework.http.client.reactive.ReactorClientHttpConnector;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import reactor.core.publisher.Flux;
import reactor.netty.http.client.HttpClient;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.net.URI;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.TimeUnit;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import static com.hjzd.ai.smartpractice.constant.ServiceConstant.*;
import static com.hjzd.ai.smartpractice.websocket.SessionManager.*;

/**
 * This example demonstrates how to connect to websocket server.
 */
public class FunasrWsClient extends WebSocketClient {

    private static final Logger logger = LoggerFactory.getLogger(FunasrWsClient.class);
    private static final Gson GSON = new GsonBuilder().setDateFormat("yyyy-MM-dd HH:mm:ss")
            .setLongSerializationPolicy(LongSerializationPolicy.STRING).disableHtmlEscaping().create();

    private boolean iseof = false;
    public static String wavPath;
    //    static String mode = "online";//2pass
    static String mode = "2pass";
    static String strChunkSize = "5,10,5";
    static int chunkInterval = 10;
    static int sendChunkSize = 1920;
    static String hotwords = "";
    static String fsthotwords = "";
    static private String srvIp = "10.252.7.232";
    static private String srvPort = "10096";
    String wavName = "javatest";
    String suffix = "wav";
    boolean isSpeaking = true;
    WebSocketSession session;
    RedisService redisService;
    private WebClientConfig webClientConfig;
    private WebClient webClient;
    private String ttsUrl;
    private String saveVoicePath;

    public class RecWavThread extends Thread {
        private FunasrWsClient funasrClient;

        public RecWavThread(FunasrWsClient funasrClient) {
            this.funasrClient = funasrClient;
        }

        public void run() {
            this.funasrClient.recWav();
        }
    }

    public FunasrWsClient(URI serverUri, Draft draft) {
        super(serverUri, draft);
    }

    public FunasrWsClient(URI serverURI, WebSocketSession session, RedisService redisService,
                          WebClientConfig webClientConfig, String ttsUrl, String saveVoicePath) {
        super(serverURI);
        this.session = session;
        this.redisService = redisService;
        this.webClientConfig = webClientConfig;
        HttpClient httpClient = HttpClient.create();
        ClientHttpConnector connector = new ReactorClientHttpConnector(httpClient);
        this.webClient = webClientConfig.webClientBuilder()
                .clientConnector(connector)
                .defaultHeader(HttpHeaders.CONTENT_TYPE, "application/json")
                .build();
        this.ttsUrl = ttsUrl;
        this.saveVoicePath = saveVoicePath;
    }

    public FunasrWsClient(URI serverUri, Map<String, String> httpHeaders) {
        super(serverUri, httpHeaders);
    }

    public void getSslContext(String keyfile, String certfile) {
        // TODO
        return;
    }

    // send json at first time
    public void sendJson() {
        try {

            JSONObject obj = new JSONObject();
            obj.put("mode", mode);
            JSONArray array = new JSONArray();
            String[] chunkList = strChunkSize.split(",");
            for (int i = 0; i < chunkList.length; i++) {
                array.add(Integer.valueOf(chunkList[i].trim()));
            }

            obj.put("chunk_size", array);
            obj.put("chunk_interval", chunkInterval);
            obj.put("wav_name", wavName);

            if (FunasrWsClient.hotwords.trim().length() > 0) {
                String regex = "\\d+";
                JSONObject jsonitems = new JSONObject();
                String[] items = FunasrWsClient.hotwords.trim().split(" ");
                Pattern pattern = Pattern.compile(regex);
                String tmpWords = "";
                for (int i = 0; i < items.length; i++) {

                    Matcher matcher = pattern.matcher(items[i]);

                    if (matcher.matches()) {

                        jsonitems.put(tmpWords.trim(), items[i].trim());
                        tmpWords = "";
                        continue;
                    }
                    tmpWords = tmpWords + items[i] + " ";

                }


                obj.put("hotwords", jsonitems.toString());
            }

            if (suffix.equals("wav")) {
                suffix = "pcm";
            }
            obj.put("wav_format", suffix);
            if (isSpeaking) {
                obj.put("is_speaking", true);
            } else {
                obj.put("is_speaking", false);
            }
            System.out.println("sendJson: " + obj);
            // return;

            send(obj.toString());

            return;
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    // send json at first time
    public void sendJson(
            String mode, String strChunkSize, int chunkInterval, String wavName, boolean isSpeaking, String suffix) {
        try {

            JSONObject obj = new JSONObject();
            obj.put("mode", mode);
            JSONArray array = new JSONArray();
            String[] chunkList = strChunkSize.split(",");
            for (int i = 0; i < chunkList.length; i++) {
                array.add(Integer.valueOf(chunkList[i].trim()));
            }

            obj.put("chunk_size", array);
            obj.put("chunk_interval", chunkInterval);
            obj.put("wav_name", wavName);

            if (FunasrWsClient.hotwords.trim().length() > 0) {
                String regex = "\\d+";
                JSONObject jsonitems = new JSONObject();
                String[] items = FunasrWsClient.hotwords.trim().split(" ");
                Pattern pattern = Pattern.compile(regex);
                String tmpWords = "";
                for (int i = 0; i < items.length; i++) {

                    Matcher matcher = pattern.matcher(items[i]);

                    if (matcher.matches()) {

                        jsonitems.put(tmpWords.trim(), items[i].trim());
                        tmpWords = "";
                        continue;
                    }
                    tmpWords = tmpWords + items[i] + " ";

                }


                obj.put("hotwords", jsonitems.toString());
            }

            if (suffix.equals("wav")) {
                suffix = "pcm";
            }
            obj.put("wav_format", suffix);
            if (isSpeaking) {
                obj.put("is_speaking", true);
            } else {
                obj.put("is_speaking", false);
            }
            System.out.println("sendJson: " + obj);
            // return;

            send(obj.toString());

            return;
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    // send json at end of wav
    public void sendEof() {
        try {
            JSONObject obj = new JSONObject();

            obj.put("is_speaking", false);

            System.out.println("sendEof: " + obj);
            // return;

            send(obj.toString());
            iseof = true;
            return;
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    // function for rec wav file
    public void sendBytes(byte[] bytes) {
        send(bytes);
    }

    // function for rec wav file
    public void recWav() {
        String fileName = FunasrWsClient.wavPath;
        String suffix = fileName.split("\\.")[fileName.split("\\.").length - 1];
        sendJson(mode, strChunkSize, chunkInterval, wavName, true, suffix);
        File file = new File(FunasrWsClient.wavPath);

        int chunkSize = sendChunkSize;
        byte[] bytes = new byte[chunkSize];

        int readSize = 0;
        try (FileInputStream fis = new FileInputStream(file)) {
            if (FunasrWsClient.wavPath.endsWith(".wav")) {
                fis.read(bytes, 0, 44); //skip first 44 wav header
            }
            readSize = fis.read(bytes, 0, chunkSize);
            while (readSize > 0) {
                // send when it is chunk size
                if (readSize == chunkSize) {
                    send(bytes); // send buf to server

                } else {
                    // send when at last or not is chunk size
                    byte[] tmpBytes = new byte[readSize];
                    for (int i = 0; i < readSize; i++) {
                        tmpBytes[i] = bytes[i];
                    }
                    send(tmpBytes);
                }
                // if not in offline mode, we simulate online stream by sleep
                if (!mode.equals("offline")) {
                    Thread.sleep(Integer.valueOf(chunkSize / 32));
                }

                readSize = fis.read(bytes, 0, chunkSize);
            }

            if (!mode.equals("offline")) {
                // if not offline, we send eof and wait for 3 seconds to close
                Thread.sleep(2000);
                sendEof();
                Thread.sleep(3000);
                close();
            } else {
                // if offline, just send eof
                sendEof();
            }

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    @Override
    public void onOpen(ServerHandshake handshakedata) {
        System.out.println("onOpen sendJson");
        sendJson();
//        RecWavThread thread = new RecWavThread(this);
//        thread.start();
    }

    @Override
    public void onMessage(String message) {
        JSONParser jsonParser = new JSONParser();
        try {
            JSONObject jsonObject = (JSONObject) jsonParser.parse(message);
            logger.info("FunasrWsClient onMessage message：{}", message);
            String mode = jsonObject.get("mode") == null ? null : jsonObject.get("mode").toString();
            Boolean isFinal = (Boolean) jsonObject.get("is_final");
            if (mode != null && mode.equals("2pass-offline")) {
                //只存储2pass修复过的文本
                String text = jsonObject.get("text").toString();
                ((StringBuilder) SessionManager.getPracticeCache(session).get(ASR_CACHE)).append(text);
            }
            if (isFinal) {
                //已是最后一句，开始文字推理。
                //1.把转写结果发送至前端
                session.sendMessage(new TextMessage(inputTextStart));
                session.sendMessage(new TextMessage(((StringBuilder) SessionManager.getPracticeCache(session).get(ASR_CACHE)).toString()));
                session.sendMessage(new TextMessage(inputTextStop));
                //2.先把转写结果存入redis
                String sessionId = SessionManager.getSession(session).getNick();
                String ChatCompletionRequestStr = redisService.get(String.format(ServiceConstant.MOCK_CHAT_UNIQUE_KEY, sessionId));
                logger.info("FunasrWsClient onMessage key:{}, ChatCompletionRequestStr：{}", String.format(ServiceConstant.MOCK_CHAT_UNIQUE_KEY, sessionId), ChatCompletionRequestStr);
                ChatCompletionRequest chatCompletionRequest;
                if (ChatCompletionRequestStr == null) {
                    throw new RuntimeException();
                } else {
                    List<ChatCompletionMessage> chatCompletionMessages = GSON.fromJson(ChatCompletionRequestStr, new TypeToken<List<ChatCompletionMessage>>() {
                    }.getType());
                    chatCompletionRequest = new ChatCompletionRequest(chatCompletionMessages, 8192, true,
                            "Qwen2-72B-Instruct");
                }
                ChatCompletionMessage chatCompletionMessage = new ChatCompletionMessage();
                chatCompletionMessage.setContent(((StringBuilder) SessionManager.getPracticeCache(session).get(ASR_CACHE)).toString());
                chatCompletionMessage.setRole(LLMRole.USER.getRole());
                chatCompletionRequest.getMessages().addLast(chatCompletionMessage);
                redisService.set(String.format(ServiceConstant.MOCK_CHAT_UNIQUE_KEY, sessionId), GSON.toJson(chatCompletionRequest.getMessages()), 7, TimeUnit.DAYS);
                //3.调用文字推理api流式返回结果
                logger.info("FunasrWsClient onMessage isFinal：{}， chatCompletionRequest: {}", isFinal, chatCompletionRequest);
                Flux<String> chatResponseFlux = webClient
                        .post()
                        .uri(INTERFERE_API_URL)
                        .contentType(MediaType.APPLICATION_JSON)
                        .bodyValue(chatCompletionRequest)
                        .retrieve()
                        .bodyToFlux(String.class);
                session.sendMessage(new TextMessage(replyTextStart));
                ChatResponseThread chatResponseThread = new ChatResponseThread(session, this.ttsUrl, this.saveVoicePath, "STEP", null, null);
                chatResponseThread.start();
                // 获取SSE结果
                chatResponseFlux.subscribe(data -> {
                            try {
                                if ("[DONE]".equals(data)) {
                                    // 保存assistant的回答
                                    String ChatCompletionResponseStr = redisService.get(String.format(ServiceConstant.MOCK_CHAT_UNIQUE_KEY, sessionId));
                                    List<ChatCompletionMessage> chatCompletionMessages;
                                    if (ChatCompletionResponseStr == null) {
                                        throw new RuntimeException();
                                    } else {
                                        chatCompletionMessages = GSON.fromJson(ChatCompletionResponseStr, new TypeToken<List<ChatCompletionMessage>>() {
                                        }.getType());
                                    }
                                    ChatCompletionMessage chatCompletionResponseMessage = new ChatCompletionMessage();
                                    chatCompletionResponseMessage.setContent(((StringBuilder) SessionManager.getPracticeCache(session).get(INTERFERE_CACHE)).toString());
                                    chatCompletionResponseMessage.setRole(LLMRole.ASSISTANT.getRole());
                                    chatCompletionMessages.addLast(chatCompletionResponseMessage);
                                    redisService.set(String.format(ServiceConstant.MOCK_CHAT_UNIQUE_KEY, sessionId), GSON.toJson(chatCompletionMessages), 7, TimeUnit.DAYS);
                                    // 向前端发送文本停止消息
                                    session.sendMessage(new TextMessage(replyTextStop));
                                    // 把本次回答的最后一部分未进入tts处理队列的文本段加入处理队列
                                    String tempTextCache = (String) SessionManager.getPracticeCache(session).get(TEMP_TEXT_CACHE);
                                    ConcurrentLinkedQueue<String> sentenceCache = (ConcurrentLinkedQueue<String>) SessionManager.getPracticeCache(session).get(SENTENCE_CACHE);
                                    if (StringUtils.isNotEmpty(tempTextCache)) {
                                        sentenceCache.add(tempTextCache);
                                    }
                                    SessionManager.getPracticeCache(session).put(TEMP_TEXT_CACHE, "");
                                    logger.info("[DONE] tempTextCache：{}, sentenceCache:{}", tempTextCache, sentenceCache);
                                    // 停止tts处理线程
                                    chatResponseThread.setStop();
                                    // 清空完整的回答缓存
                                    ((StringBuilder) SessionManager.getPracticeCache(session).get(INTERFERE_CACHE)).setLength(0);
                                } else {
                                    ChatCompletionResponse response = GSON.fromJson(data, ChatCompletionResponse.class);
                                    if (response != null && response.getChoices() != null
                                            && response.getChoices().getFirst().getDelta() != null
                                            && response.getChoices().getFirst().getDelta().getContent() != null) {
                                        String content = response.getChoices().getFirst().getDelta().getContent().replace("\"", "");
                                        session.sendMessage(new TextMessage(content));
                                        StringBuilder interfereCache = (StringBuilder) SessionManager.getPracticeCache(session).get(INTERFERE_CACHE);
                                        interfereCache.append(content);
                                        String tempTextCache = (String) SessionManager.getPracticeCache(session).get(TEMP_TEXT_CACHE);
                                        tempTextCache = tempTextCache + content;
                                        ConcurrentLinkedQueue<String> textQueue =
                                                (ConcurrentLinkedQueue<String>) SessionManager.getPracticeCache(session).get(SENTENCE_CACHE);
                                        logger.info("[NOT DONE] tempTextCache：{}", tempTextCache);
                                        tempTextCache = sentenceSplit(SENTENCE_SPLIT_REGEX, tempTextCache, textQueue);
                                        logger.info("[NOT DONE] sentenceCache：{}, tempTextCache:{}", textQueue, tempTextCache);
                                        SessionManager.getPracticeCache(session).put(TEMP_TEXT_CACHE, tempTextCache);
                                        logger.info("subscribe content:{}, interfereCache:{}, textQueue:{}", content, interfereCache, textQueue);
                                    }
                                }
                            } catch (Exception e) {
                                logger.error("subscribe data error,", e);
                            }
                        },
                        error -> {
                            logger.error("subscribe error,", error);
                        }
                );
                logger.info("文字推理结束");
            }
        } catch (ParseException | IOException e) {
            logger.error("session sendMessage error,", e);
        }
    }

    @Override
    public void onClose(int code, String reason, boolean remote) {

        System.out.println(
                "Connection closed by "
                        + (remote ? "remote peer" : "us")
                        + " Code: "
                        + code
                        + " Reason: "
                        + reason);
    }

    @Override
    public void onError(Exception ex) {
        System.out.println("ex: " + ex);
        ex.printStackTrace();
        // if the error is fatal then onClose will be called additionally
    }

    private static String sentenceSplit(String regex, String text, ConcurrentLinkedQueue<String> queue) {
        // 创建Pattern对象
        Pattern pattern = Pattern.compile(regex);

        // 创建matcher对象
        Matcher matcher = pattern.matcher(text);

        // 使用Matcher对象的find方法分割字符串
        int lastScanEnd = 0;
        int lastAddedEnd = 0;
        StringBuilder lastSplit = new StringBuilder();
        while (matcher.find()) {
            String word = text.substring(lastScanEnd, matcher.end());
            lastSplit.append(word);
            lastScanEnd = matcher.end();
            if (lastSplit.length() > SENTENCE_SEND_MIN_WORDS_LIMIT) {
                queue.add(lastSplit.toString());
                lastSplit.setLength(0);
                lastAddedEnd = matcher.end();
            }
        }
        // 返回未入队列的字符
        return text.substring(lastAddedEnd);
    }
//    public static void main(String[] args) throws URISyntaxException {
//        ArgumentParser parser = ArgumentParsers.newArgumentParser("ws client").defaultHelp(true);
//        parser
//                .addArgument("--port")
//                .help("Port on which to listen.")
//                .setDefault("8889")
//                .type(String.class)
//                .required(false);
//        parser
//                .addArgument("--host")
//                .help("the IP address of server.")
//                .setDefault("127.0.0.1")
//                .type(String.class)
//                .required(false);
//        parser
//                .addArgument("--audio_in")
//                .help("wav path for decoding.")
//                .setDefault("asr_example.wav")
//                .type(String.class)
//                .required(false);
//        parser
//                .addArgument("--num_threads")
//                .help("num of threads for test.")
//                .setDefault(1)
//                .type(Integer.class)
//                .required(false);
//        parser
//                .addArgument("--chunk_size")
//                .help("chunk size for asr.")
//                .setDefault("5, 10, 5")
//                .type(String.class)
//                .required(false);
//        parser
//                .addArgument("--chunk_interval")
//                .help("chunk for asr.")
//                .setDefault(10)
//                .type(Integer.class)
//                .required(false);
//
//        parser
//                .addArgument("--mode")
//                .help("mode for asr.")
//                .setDefault("offline")
//                .type(String.class)
//                .required(false);
//        parser
//                .addArgument("--hotwords")
//                .help("hotwords, splited by space, hello 30 nihao 40")
//                .setDefault("")
//                .type(String.class)
//                .required(false);
//        String srvIp = "";
//        String srvPort = "";
//        String wavPath = "";
//        int numThreads = 1;
//        String chunk_size = "";
//        int chunk_interval = 10;
//        String strmode = "offline";
//        String hot = "";
//        try {
//            Namespace ns = parser.parseArgs(args);
//            srvIp = ns.get("host");
//            srvPort = ns.get("port");
//            wavPath = ns.get("audio_in");
//            numThreads = ns.get("num_threads");
//            chunk_size = ns.get("chunk_size");
//            chunk_interval = ns.get("chunk_interval");
//            strmode = ns.get("mode");
//            hot = ns.get("hotwords");
//            System.out.println(srvPort);
//
//        } catch (ArgumentParserException ex) {
//            ex.getParser().handleError(ex);
//            return;
//        }
//
//        FunasrWsClient.strChunkSize = chunk_size;
//        FunasrWsClient.chunkInterval = chunk_interval;
//        FunasrWsClient.wavPath = wavPath;
//        FunasrWsClient.mode = strmode;
//        FunasrWsClient.hotwords = hot;
//        System.out.println(
//                "serIp="
//                        + srvIp
//                        + ",srvPort="
//                        + srvPort
//                        + ",wavPath="
//                        + wavPath
//                        + ",strChunkSize"
//                        + strChunkSize);
//
//        class ClientThread implements Runnable {
//
//            String srvIp;
//            String srvPort;
//
//            ClientThread(String srvIp, String srvPort, String wavPath) {
//                this.srvIp = srvIp;
//                this.srvPort = srvPort;
//            }
//
//            public void run() {
//                try {
//
//                    int RATE = 16000;
//                    String[] chunkList = strChunkSize.split(",");
//                    int int_chunk_size = 60 * Integer.valueOf(chunkList[1].trim()) / chunkInterval;
//                    int CHUNK = RATE / 1000 * int_chunk_size;
//                    int stride = 60 * Integer.parseInt(chunkList[1].trim()) / chunkInterval / 1000 * 16000 * 2;
//                    System.out.println("chunk_size:" + String.valueOf(int_chunk_size));
//                    System.out.println("CHUNK:" + CHUNK);
//                    System.out.println("stride:" + String.valueOf(stride));
//                    FunasrWsClient.sendChunkSize = CHUNK * 2;
//                    String wsAddress = "ws://" + srvIp + ":" + srvPort;
//                    FunasrWsClient c = new FunasrWsClient(new URI(wsAddress));
//                    c.connect();
//                    System.out.println("wsAddress:" + wsAddress);
//                } catch (Exception e) {
//                    e.printStackTrace();
//                    System.out.println("e:" + e);
//                }
//            }
//        }
//        for (int i = 0; i < numThreads; i++) {
//            System.out.println("Thread1 is running...");
//            Thread t = new Thread(new ClientThread(srvIp, srvPort, wavPath));
//            t.start();
//        }
//    }
}
